'''
Created on Jul 28, 2010

@author: oabalbin
'''

import numpy as np
import rpy2.robjects as robjects
import rpy2.robjects.numpy2ri
from rpy2.robjects.packages import importr

from optparse import OptionParser
from collections import deque, defaultdict

def list_of_names_in_line_dictionary(inputfile):     

    genlist=defaultdict()
    for line in inputfile:
        if line.startswith('#'):
            continue        
        
        line = line.strip('\n')
        genlist0 = line.split('\t')
        #print genlist0[0], genlist0[1:]
        genlist[genlist0[0]]=genlist0[2:]
    
    return genlist    

def list_of_names_in_line(inputfile):     

    genlist=deque()
    for line in inputfile:
        if line.startswith('#'):
            continue        
        line = line.strip('\n')
        genlist = line.split('\t')
            
    return genlist


def list_of_names(inputfile):
    
    genlist=deque()
    for line in inputfile:
        if line.startswith('#'):
            continue
                
        line = line.strip('\n')
        fields = line.split('\t')
        #print fields 
        # To skip headers, star reading samples in column 7 of the file
        if fields[0][0] == '#': 
            continue                             
        if fields[0]=='':
            continue
        else:
            'Two first fields are reserved for the name of the gene set and description'
            for gen in fields[2:]:
                genlist.append(gen)
    
    return genlist


def values_for_enrichment_test(mygeneset, universe_geneset, external_geneset):
    """
    x, q     vector of quantiles representing the number of white balls drawn without replacement from an urn which contains both black and white balls.
    m     the number of white balls in the urn.
    n     the number of black balls in the urn.
    k     the number of balls drawn from the urn.
    """
    notest=[]
    
    drawn_white_balls=len(set(mygeneset).intersection(set(external_geneset)))
    if drawn_white_balls ==0:
        return notest
    
    total_white_balls_urn=len(set(external_geneset).intersection(set(universe_geneset)))
    total_black_balls_urn= len(set(universe_geneset)) - total_white_balls_urn #len(set(universe_geneset).difference(set(external_geneset)))
    total_balls_drawn=len(mygeneset)
    
    values_enrichment_test = [drawn_white_balls, total_white_balls_urn, total_black_balls_urn, total_balls_drawn]
    
    return values_enrichment_test


def calc_hypergeometric_test(cell_values):
    """
    It calculates an hypergeometric test using the parameters given
    phyper(q, m, n, k, lower.tail = TRUE, log.p = FALSE)
    Arguments
    x, q     vector of quantiles representing the number of white balls drawn without replacement from an urn which contains both black and white balls.
    m     the number of white balls in the urn.
    n     the number of black balls in the urn.
    k     the number of balls drawn from the urn.
    log, log.p     logical; if TRUE, probabilities p are given as log(p).
    lower.tail     logical; if TRUE (default), probabilities a    
    """
    
    r = robjects.r
    r("library(stats)")
    r("library(qvalue)")
    rphyper = r["phyper"]
    drawn_white_balls, total_white_balls, total_black_balls, total_balls_drawn = \
    cell_values[0],cell_values[1],cell_values[2],cell_values[3]
    
    prob_enriched = rphyper(drawn_white_balls, total_white_balls, total_black_balls, total_balls_drawn, lower_tail=False);
        
    return np.array(prob_enriched)[0]


def calculate_FDR(pvalues_array):
    """
    It returns a FDR for the array of pvalues
    
    # Start by sorting all rows in increasing order of t-test p-values for the brain (adult vs. fetal) analysis.
    # Note the number of probesets we're analyzing
    # Copy the last raw p-value into the column of FDR p-values.
    # For the fields above the last, apply the formula
    fdr = min(raw * (num_rows/rank_this_probeset), fdr_for_gene_one_row_below)
    ex: =MIN(B7923 * 7923/RANK(B7923,B$2:B7924,1),F7924)
    Note that the rank is calculated for the list of raw p-values in ascending order.
    # Paste this formula in the cells above to correct for all probesets.
    """
    
    number_of_test=len(pvalues_array)
    pvalues_array2sort = np.copy(pvalues_array)
    inidicators = np.array(range(number_of_test))
    sort_indicators = np.argsort(pvalues_array2sort)  
    qvalue_array = np.empty((number_of_test,2))
    qvalues=np.empty(number_of_test)
    # sort the original pvalues
    pvalues_array2sort.sort()
    
    First=True
    for i in inidicators[::-1]:
        rank=i+1
        k=sort_indicators[i]
        if First:
            qvalue_array[k,0] = pvalues_array2sort[i]
            qvalue_array[k,1] = pvalues_array2sort[i]
            qvalues[i] = pvalues_array2sort[i]
            First=False
        else:
            qvalue_array[k,0] = pvalues_array2sort[i]
            pvalue=pvalues_array2sort[i]
            #print pvalue, pvalue*(number_of_test/float(rank)), qvalue_array[i+1,1], np.minimum( pvalue*(number_of_test/float(rank)), qvalue_array[i+1,1] )
            qvalues[i] = np.minimum( pvalue*(number_of_test/float(rank)), qvalues[i+1] )
            qvalue_array[k,1] = qvalues[i]
    
    #qvalue_array = qvalue_array[sort_indicators,:]
    return qvalue_array#, sort_indicators


def calc_qvalue(pvalues, FDR):
    
    r = robjects.r
    r("library(qvalue)")
    rqvalue=r["qvalue"]
    qvalues = rqvalue(pvalues,fdr_level=FDR)
    
    return qvalues
    
if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--mygenset", dest="mygenset",
                            help="mygenset file to use")
    optionparser.add_option("-u", "--universe_geneset", dest="universe_geneset",
                            help="universe_geneset file to use")
    optionparser.add_option("-s", "--external_sets", dest="external_sets",
                            help="external_sets file to use")
    optionparser.add_option("-o", "--outfile", dest="outfile",
                            help="external_sets file to use")

   
    (options, args) = optionparser.parse_args()

    mygeneset=list_of_names(open(options.mygenset))
    universe_geneset=list_of_names(open(options.universe_geneset))
    external_sets=list_of_names_in_line_dictionary(open(options.external_sets))
    fdr_treshold = 0.5

    #print mygeneset
    #print universe_geneset
    
    outfile=open(options.outfile,'w')
    pvalues_list, performed_tests, test_values=deque(), deque(), deque()
    
    for set_name, gene_set  in external_sets.iteritems():
        values_enrichment_test = values_for_enrichment_test(mygeneset, universe_geneset, gene_set)
        if not values_enrichment_test:
            continue
        
        pvalue = calc_hypergeometric_test(values_enrichment_test)
        performed_tests.append((set_name,values_enrichment_test,pvalue)) 
        pvalues_list.append(pvalue)
        
        if pvalue < 0.01:
            print set_name, pvalue
    
    #pvalues_list = range(len(performed_tests))
    fdr_values = calculate_FDR(np.array(pvalues_list))

    for i, thisttest in enumerate(performed_tests):
        if fdr_values[i][1] < fdr_treshold:
            print thisttest, str(fdr_values[i][0]), str(fdr_values[i][1])
            line=[thisttest[0],str(fdr_values[i][0]), str(fdr_values[i][1])]
            outfile.write(",".join(line).replace(',','\t')+'\t test_values \t'+",".join(map(str,thisttest[1])).replace(',','\t')+'\n')

    
    
        